从一个朴素的问题说起
刚开始做 NLP 训练的时候，我的数据管线极其朴素：pandas 读一个 CSV，tokenizer 过一遍，塞进 PyTorch 的 TensorDataset，完事。数据量小、单机跑实验，这套流程完全够用。
但问题很快来了。数据量从几十万涨到几千万行的时候，pandas 的 read_csv 开始把内存吃满；我想对数据做一些列变换（比如把 label 从字符串映射成整数），pandas 的 apply 慢得让人抓狂；多个数据集需要按 split 管理（train / validation / test），我自己维护字典，越写越乱；最要命的是，我想换一个数据集试试效果，整个读取和预处理流程得重写一遍。
datasets 包就是在这个背景下被 HuggingFace 团队做出来的。它表面上是一个”加载数据集的工具”，但如果你只用 load_dataset 这一个函数，其实只用了它大概 20% 的能力。真正有意思的是它底层的数据表示方式和转换管线的设计。
底层为什么选了 Apache Arrow
datasets 包最核心的一个设计决策，是底层数据存储用了 Apache Arrow 格式，而不是 Python 原生的 list/dict，也不是 NumPy 数组，更不是 pandas DataFrame。
这个选择不是拍脑袋的。Arrow 是一种列式内存格式，你可以把它理解成一种”内存里的 Parquet”。它带来几个对 ML 训练管线极其关键的特性：
第一，零拷贝读取。 Arrow 的数据布局是连续的、固定 schema 的内存块。当 datasets 把数据交给 PyTorch 或者 TensorFlow 的时候，不需要把 Python 对象逐个序列化再反序列化，可以直接把内存指针交过去。这在 batch size 很大、序列很长的场景下，省掉的拷贝开销非常可观。
第二，内存映射（memory-mapped）。 数据集可以存在磁盘上，用 mmap 的方式映射到虚拟内存空间。操作系统按需分页加载，你不需要把整个几十 GB 的数据集一次性读进物理内存。datasets 默认就会把下载好的数据集缓存成 Arrow 文件，后续加载走的就是 mmap。
第三，列式访问天然高效。 训练时你通常只需要某几列（比如 input_ids 和 labels），列式存储意味着你只读需要的列，不需要的列根本不碰。如果是行式存储，每一行都得完整读出来再丢弃不需要的字段。
比如在处理一个 20GB 的文本数据集时，pandas 直接 read_csv 需要 20GB 以上的内存，而 datasets 加载之后，nvidia-smi 和 htop 显示实际物理内存占用只有几个 GB——因为 mmap 只把当前访问到的页加载进来了。
三个核心抽象：Dataset、DatasetDict、IterableDataset
datasets 包的数据模型其实就三层。
Dataset 是最基本的单元，你可以把它理解成一张”内存中的表”。每一行是一个样本，每一列是一个字段（text、label、input_ids……）。它支持索引、切片、迭代，接口设计得很像 Python list，但底层是 Arrow Table。
DatasetDict 是一个 str -> Dataset 的映射，用来组织同一个数据集的不同 split。load_dataset("imdb") 返回的就是一个 DatasetDict，里面有 "train" 和 "test" 两个 key。这个设计看起来平平无奇，但它让你对整个数据集的操作（比如 tokenize）可以统一地作用在所有 split 上，不用自己写循环。
IterableDataset 是流式场景的抽象。它不把所有数据加载到内存（甚至不下载到本地），而是像 Python generator 一样，你迭代一次它吐一条。这在数据集大到磁盘都装不下、或者数据在网络流上实时产生的场景下是刚需。
这三层抽象覆盖了从”小数据集全量加载”到”超大数据集流式消费”的完整谱系，设计上是自洽的。
map：整个包最精华的设计
如果说 datasets 包只能保留一个功能，我会选 map。
Dataset.map() 的语义很简单：对数据集中的每一行（或每一个 batch）应用一个函数，返回一个新的 Dataset。但它的实现细节和工程考量，才是真正有意思的地方。
批处理与并行。 你可以设置 batched=True 让函数一次接收一批数据（而不是逐条调用），这对 tokenizer 这种本身支持批量处理的函数来说，速度差异是数量级的。再配合 num_proc 参数，datasets 会用多进程并行地跑 map，每个进程处理数据的一个 shard。比如8 核机器上对 500 万条数据做 tokenize，num_proc=8 比单进程理论上快将近 6 倍（没到 8 倍是因为有 IO 和进程通信开销）。
惰性与缓存。 map 的结果会被缓存到磁盘（默认在 ~/.cache/huggingface/datasets/ 下）。如果你用完全相同的函数和参数再跑一次 map，它直接从缓存读，不会重复计算。这在调试训练代码的时候特别有用——你改了模型结构想重跑，数据预处理这一步直接跳过。缓存的 key 是根据函数签名和参数算出来的 hash，所以如果你改了函数内部的逻辑，缓存会自动失效。
列的增删。 map 默认会把原列和新列都保留。但很多时候你 tokenize 完之后就不需要原始的 text 列了，可以用 remove_columns 参数在 map 的同时把不需要的列丢掉，避免后续加载时白白多读数据。
一个典型的训练前数据准备大概长这样：
from datasets import load_dataset
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
ds = load_dataset("ag_news")
def tokenize_fn(examples):
return tokenizer(examples["text"], truncation=True, max_length=128)
tokenized_ds = ds.map(tokenize_fn, batched=True, num_proc=4, remove_columns=["text"])这段代码背后，datasets 帮你做了分片、多进程调度、Arrow 格式写入、缓存管理、列裁剪。你自己不需要碰任何这些脏活。
接入 PyTorch 训练：没那么简单但也没那么复杂
datasets 包原生支持把 Dataset 的格式切换成 PyTorch 张量。调用 dataset.set_format("torch", columns=["input_ids", "attention_mask", "labels"]) 之后，迭代出来的每个样本就是 torch.Tensor 了。然后你可以直接把它丢给 torch.utils.data.DataLoader：
from torch.utils.data import DataLoader
train_loader = DataLoader(tokenized_ds["train"], batch_size=32, shuffle=True)这里有一个容易踩的坑：datasets 的 Dataset 实现了 __getitem__，但它的 shuffle 和 DataLoader 的 shuffle 是两回事。 如果你用 IterableDataset（流式模式），DataLoader 的 shuffle=True 是无效的，因为 IterableDataset 不支持随机访问。你需要在数据准备阶段自己调用 dataset.shuffle(seed=42, buffer_size=10000) 来做 buffer shuffle。这个 buffer_size 的含义是从流中先读入 10000 条到内存，在这个 buffer 内部打乱顺序再吐出，它是一个近似 shuffle，buffer 越大越接近全局 shuffle，但内存开销也越大。
另一个实际中经常碰到的问题是 DataLoader 的 num_workers 和 datasets 的 mmap 之间的交互。mmap 的文件描述符在 fork 子进程时会被继承，大部分情况下这没问题，但如果你用的是非常老的 Linux 内核或者某些特殊的文件系统，可能会碰到 fd 泄漏或者 SIGBUS。遇到这种情况，一个务实的 workaround 是把 num_workers 设成 0，或者在 DataLoader 里用 multiprocessing_context="spawn" 代替默认的 fork。
流式处理：数据大到装不下怎么办
当数据集到了几百 GB 甚至 TB 级别（比如 The Pile、RedPajama、FineWeb），全量下载再加载是不现实的。datasets 提供了 streaming=True 模式：
ds = load_dataset("allenai/c4", "en", streaming=True, split="train")
for sample in ds:
# 逐条处理，内存里只有一条数据
...流式模式下返回的是 IterableDataset，它本质上是一个包装了 HTTP 请求的 generator。数据按需从 HuggingFace Hub（或者你指定的本地/远程路径）拉取，不会落盘缓存整个数据集。
但流式模式有一个根本性的限制：你无法随机访问。 不能做 dataset[42]，不能做全局 shuffle，不能做需要两遍扫描的操作（比如计算全局均值来做归一化）。如果你的训练流程需要这些，要么先用流式模式把数据预处理成 Arrow 文件存到磁盘，要么接受近似（buffer shuffle、在线归一化）。
在实际的大规模预训练里，我见过的做法是：用流式模式做第一遍 tokenize，把结果写成本地的 Arrow/Parquet 分片文件；然后训练时用 load_dataset("arrow", data_files=[...]) 加载这些本地文件，这时候就回到了支持随机访问的模式。这相当于把”下载 + 预处理”和”训练消费”解耦成了两个阶段。
自定义数据集：别总想着写 JSON
datasets 对自定义数据集的支持非常灵活。最简单的方式是直接从一个字典构建：
from datasets import Dataset
data = {"text": ["hello", "world"], "label": [0, 1]}
ds = Dataset.from_dict(data)稍微复杂一点，你可以从本地文件加载，支持 CSV、JSON、JSONL、Parquet、Arrow、文本文件等格式：
ds = load_dataset("json", data_files={"train": "train.jsonl", "test": "test.jsonl"})但如果你有一个很特殊的格式（比如每条样本是一个目录下的图片加一个 XML 标注），你需要写一个自定义的 GeneratorBasedBuilder，实现 _generate_examples 方法。这个接口的设计思路是：你只需要写一个 Python generator，yield 出 (key, example_dict) 对，datasets 会帮你处理 Arrow 序列化、缓存、split 管理。
我自己的经验是，能用 Parquet 就用 Parquet。Parquet 是列式存储、自带压缩、schema 信息嵌在文件头里，datasets 读 Parquet 的速度比读 JSONL 快很多，而且天然支持列裁剪。如果你的原始数据是 JSONL，值得花几分钟先转成 Parquet 再喂给 datasets。
几个工程上真实踩过的坑
缓存路径的磁盘炸弹。 每次 map 操作都会生成新的缓存文件。如果你反复调试预处理逻辑，~/.cache/huggingface/datasets/ 会迅速膨胀到几十甚至上百 GB。设置 HF_DATASETS_CACHE 环境变量把缓存指到大容量磁盘上，或者定期手动清理，是生产环境的基本操作。
map 里用了闭包或者 lambda 导致缓存失效。 datasets 的缓存 key 依赖函数的可序列化性。如果你 map 的函数引用了外部变量（闭包），或者用了 lambda，有时候缓存 hash 会不稳定，导致每次运行都重新计算。最稳妥的做法是把 map 函数定义成顶层函数，需要的参数通过 fn_kwargs 传进去。
多进程 map 和 GPU 操作冲突。 如果你 map 的函数里涉及 CUDA 操作（比如用一个 GPU 上的模型做 embedding），num_proc > 1 会出问题，因为 CUDA context 不能被 fork。这种情况要么 num_proc=1，要么用 multiprocessing 的 spawn 模式，要么干脆把 GPU 推理和 datasets 的 map 分成两个阶段。
Dataset 的 select 和 filter 的性能差异。 select 是基于索引的，只是创建了一个新的索引映射，几乎零开销。filter 需要对每一行执行谓词函数，是 O(n) 的。如果你只是想取一个子集，优先用 select(range(0, len(ds), 10)) 这种方式，而不是 filter(lambda x: x["id"] % 10 == 0)。
最后说两句
datasets 包的设计哲学，我觉得可以概括成一句话：把数据当成一个”有 schema 的、可变换的、可缓存的、格式统一的”对象来管理，而不是一堆散落在内存里的 Python list。
它不完美。调试的时候 Arrow 格式的报错信息有时候不太友好；流式模式的限制在真实场景下确实会碰到；和 PyTorch DataLoader 的衔接偶尔需要一些胶水代码。但在”从原始数据到训练 tensor”这条管线上，它目前是我用过的工程完成度最高的方案。它把那些你不想关心但不得不做的脏活——格式转换、缓存、并行、内存管理——都封装在了一个还算优雅的 API 后面，让你能把精力放在真正重要的事情上：数据本身和模型。
如果你的训练数据还停留在”pandas 读进来，for 循环处理，手动切 split”的阶段，真的建议花一个下午把 datasets 的文档过一遍。不一定是为了用 HuggingFace Hub 上的公开数据集，光是它本地的 Arrow 缓存、多进程 map、和 DataLoader 的无缝对接，就足够让你的数据管线干净一个量级。